import os
import os.path as osp

from PIL import Image
from torch.utils.data import Dataset
from torchvision import transforms

TARGET_DATA_DIR = "../../fewshot_cd_baselines/target_domain"
DATASET_DICT = {
    "cars": "cars/novel",
    "cropdisease": "CropDisease",
    "cub": "CUB/novel",
    "eurosat": "EuroSAT",
    "isic": "ISIC",
    "places": "Places/novel",
    "plantae": "Plantae/novel",
    "chestx": "XChest",
}


class GeneralLoader(Dataset):
    def __init__(self, setname, args, train_augmentation=None):
        assert args.dataset in DATASET_DICT, "unknown dataset"

        data_path = osp.join(TARGET_DATA_DIR, DATASET_DICT[args.dataset])

        # if setname == "train":
        #     THE_PATH = osp.join(DATASET_DIR, "train")
        #     label_list = os.listdir(THE_PATH)
        # elif setname == "test":
        #     THE_PATH = osp.join(DATASET_DIR, "test")
        #     label_list = os.listdir(THE_PATH)
        # elif setname == "val":
        #     THE_PATH = osp.join(DATASET_DIR, "val")
        #     label_list = os.listdir(THE_PATH)
        # else:
        #     raise ValueError("Incorrect set name. Please check!")
        label_list = os.listdir(data_path)

        data = []
        label = []

        folders = [
            osp.join(data_path, label)
            for label in label_list
            if osp.isdir(osp.join(data_path, label))
        ]

        for idx, this_folder in enumerate(folders):
            this_folder_images = os.listdir(this_folder)
            for image_path in this_folder_images:
                data.append(osp.join(this_folder, image_path))
                label.append(idx)

        self.data = data
        self.label = label
        self.num_class = len(folders)  # len(set(label))

        # Transformation
        if setname == "train" and train_augmentation is not None:
            self.transform = train_augmentation

        elif (setname == "val" or setname == "test") and train_augmentation is None:
            image_size = args.image_size
            if image_size == 224:
                img_resize = 256
            elif image_size == 84:
                img_resize = 92
            else:
                ValueError("Image size not supported at the moment.")
            self.transform = transforms.Compose(
                [
                    transforms.Resize([img_resize, img_resize]),
                    transforms.CenterCrop(image_size),
                    transforms.ToTensor(),
                    transforms.Normalize(
                        (0.485, 0.456, 0.406), (0.229, 0.224, 0.225)
                    ),  # ImageNet standard
                ]
            )
        else:
            ValueError("Set name or train augmentation corrupt. Please check!")

    def __len__(self):
        return len(self.data)

    def __getitem__(self, i):
        path, label = self.data[i], self.label[i]
        image = self.transform(Image.open(path).convert("RGB"))
        return image, label


if __name__ == "__main__":
    pass
